#!/usr/bin/env python3
"""
从parquet文件中提取top_5为0的数据并导出为JSONL格式。

主要功能：
- 查找指定目录下的所有parquet文件
- 过滤出top_5指标为0的样本数据
- 将过滤后的数据导出为JSONL格式文件
- 提供详细的统计信息

数据格式要求：
- 输入parquet必须包含'metrics'列或'info.metrics'字段
- 支持多种top_5指标格式（带任意前缀）

使用示例：
  # 从指定文件夹提取top_5为0的数据
  python extract_low_top5_data.py --input /path/to/folder --output /path/to/output.jsonl
"""

import os
import sys
import json
import argparse
import logging
from datetime import datetime
from typing import List, Dict, Optional, Any

import pandas as pd
import numpy as np


class NumpyEncoder(json.JSONEncoder):
    """自定义JSON编码器，用于处理numpy数据类型"""
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.bool_):
            return bool(obj)
        return super().default(obj)


def configure_logging(level: str = "INFO") -> None:
    """配置日志输出"""
    root = logging.getLogger()
    root.handlers.clear()
    root.setLevel(logging.getLevelName(level.upper()))
    handler = logging.StreamHandler(stream=sys.stdout)
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    handler.setFormatter(formatter)
    root.addHandler(handler)


def read_parquet_safe(path: str) -> pd.DataFrame:
    """安全读取parquet文件"""
    if not os.path.isfile(path):
        return pd.DataFrame()
    try:
        return pd.read_parquet(path, engine="pyarrow")
    except Exception:
        return pd.read_parquet(path, engine="fastparquet")


def find_top5_metric(metrics: Dict[str, Any]) -> Optional[tuple]:
    """
    查找top_5相关指标

    Args:
        metrics: metrics字典

    Returns:
        (top5_key, top5_value) 如果找到，否则返回None
    """
    for key, value in metrics.items():
        if "top_5" in key.lower():
            return key, float(value)
    return None


def extract_low_top5_samples(df: pd.DataFrame) -> List[Dict[str, Any]]:
    """
    从DataFrame中提取top_5为0的样本

    Args:
        df: 输入DataFrame

    Returns:
        过滤后的样本列表
    """
    low_top5_samples = []

    # 检查metrics可能在顶层列或info列中
    has_metrics_column = "metrics" in df.columns
    has_info_column = "info" in df.columns

    for idx, row in df.iterrows():
        sample_dict = row.to_dict()

        # 获取metrics数据
        metrics = {}
        if has_metrics_column:
            metrics = row.get("metrics", {})
        elif has_info_column:
            info = row.get("info", {})
            if isinstance(info, dict):
                metrics = info.get("metrics", {})

        if not metrics:
            continue

        # 查找top_5指标
        top5_info = find_top5_metric(metrics)
        if top5_info is None:
            continue

        top5_key, top5_value = top5_info
        if top5_value == 0.0:
            # 记录完整的样本信息
            sample_copy = sample_dict.copy()
            sample_copy["_metadata"] = {
                "top5_metric_key": top5_key,
                "top5_value": top5_value,
                "row_index": idx
            }
            low_top5_samples.append(sample_copy)

    return low_top5_samples


def process_directory(input_dir: str, output_file: str) -> Dict[str, Any]:
    """
    处理指定目录下的所有parquet文件

    Args:
        input_dir: 输入目录
        output_file: 输出JSONL文件

    Returns:
        处理统计信息
    """
    if not os.path.isdir(input_dir):
        raise ValueError(f"输入目录不存在: {input_dir}")

    # 查找所有parquet文件
    parquet_files = []
    for filename in os.listdir(input_dir):
        if filename.endswith('.parquet'):
            file_path = os.path.join(input_dir, filename)
            parquet_files.append(file_path)

    if not parquet_files:
        raise ValueError(f"目录中未找到parquet文件: {input_dir}")

    logging.info(f"找到 {len(parquet_files)} 个parquet文件")

    all_samples = []
    total_processed = 0
    total_filtered = 0

    # 处理每个parquet文件
    for file_path in parquet_files:
        try:
            df = read_parquet_safe(file_path)
            if df.empty:
                logging.warning(f"文件为空: {file_path}")
                continue

            logging.info(f"处理文件: {file_path} ({len(df)} 行)")

            # 提取top_5为0的样本
            low_top5_samples = extract_low_top5_samples(df)

            all_samples.extend(low_top5_samples)
            total_processed += len(df)
            total_filtered += len(low_top5_samples)

            if low_top5_samples:
                logging.info(f"  ✓ 找到 {len(low_top5_samples)} 个top_5=0的样本")

        except Exception as e:
            logging.error(f"处理文件失败 {file_path}: {e}")
            continue

    # 写入JSONL文件
    if all_samples:
        with open(output_file, "w", encoding="utf-8") as f:
            for sample in all_samples:
                json.dump(sample, f, ensure_ascii=False, cls=NumpyEncoder)
                f.write("\n")

        logging.info(f"✓ 成功导出 {len(all_samples)} 个样本到: {output_file}")
    else:
        logging.warning("未找到任何top_5=0的样本")

    # 生成统计信息
    stats = {
        "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
        "input_directory": input_dir,
        "total_parquet_files": len(parquet_files),
        "total_processed_rows": total_processed,
        "total_low_top5_samples": len(all_samples),
        "output_file": output_file,
        "filter_percentage": (len(all_samples) / total_processed * 100) if total_processed > 0 else 0
    }

    return stats


def print_statistics(stats: Dict[str, Any]) -> None:
    """打印处理统计信息"""
    print("\n" + "=" * 60)
    print("📊 处理统计信息")
    print("=" * 60)
    print(f"📂 输入目录: {stats['input_directory']}")
    print(f"📁 处理文件数: {stats['total_parquet_files']}")
    print(f"📊 总处理行数: {stats['total_processed_rows']}")
    print(f"🎯 找到top_5=0的样本: {stats['total_low_top5_samples']}")
    print(f"📈 过滤比例: {stats['filter_percentage']:.2f}%")
    print(f"💾 输出文件: {stats['output_file']}")
    print(f"⏰ 处理时间: {stats['timestamp']}")
    print("=" * 60)


def parse_args() -> argparse.Namespace:
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="从parquet文件中提取top_5为0的数据并导出为JSONL格式",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  # 从指定文件夹提取
  python extract_low_top5_data.py --input /path/to/folder --output output.jsonl

  # 指定日志级别
  python extract_low_top5_data.py --input /path/to/folder --output output.jsonl --log-level DEBUG
        """
    )

    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="输入目录路径，包含parquet文件"
    )

    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="输出JSONL文件路径"
    )

    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="日志级别（默认: INFO）"
    )

    return parser.parse_args()


def main() -> None:
    """主函数"""
    args = parse_args()
    configure_logging(args.log_level)

    # 打印配置信息
    print("=" * 60)
    print("🔍 top_5=0数据提取工具")
    print("=" * 60)
    print(f"📂 输入目录: {args.input}")
    print(f"💾 输出文件: {args.output}")
    print(f"📊 日志级别: {args.log_level}")
    print("=" * 60)
    print()

    try:
        # 处理数据
        stats = process_directory(args.input, args.output)

        # 打印统计信息
        print_statistics(stats)

        print("\n✅ 处理完成!")

    except Exception as e:
        logging.exception("处理过程中发生错误")
        raise SystemExit(f"❌ 错误: {e}")


if __name__ == "__main__":
    main()